import threading
from collections.abc import Callable
from functools import partial
from typing import ParamSpec, TypeVar, Union, overload

P = ParamSpec('P')
R = TypeVar('R')


@overload
def locked(lock: Union[threading.Lock, threading.RLock, None] = None) -> Callable[[Callable[P, R]], Callable[P, R]]:
    pass


@overload
def locked(lock: Union[threading.Lock, threading.RLock, None], func: Callable[P, R]) -> Callable[P, R]:
    pass


def locked(lock: Union[threading.Lock, threading.RLock, None] = None, func: Callable[P, R] | None = None):
    if lock is None:
        lock = threading.RLock()

    if func is None:
        if isinstance(lock, Callable):
            return locked(None, lock)
        return partial(locked, lock)

    def locked_func(*args: P.args, **kwargs: P.kwargs) -> R:
        with lock:
            return func(*args, **kwargs)

    return locked_func
